# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import logging
import time
import traceback

from coin.base.timestamp import get_timestamp
from coin.support.accounting.logic.key_util import get_key_file
from coin.exchange.kr_rest.feed.feed_impl import (
    generate_public_client,)
from coin.exchange.kr_rest.order.order_impl import (
    generate_private_client,)
from coin.proto.coin_query_pb2 import (
    ProductOrderElement,)
from coin.proto.exchange_status_pb2 import (
    OrderAttemptProto,
    RestStatusProto,
)
from coin.exchange.base.order_enum import (
    get_str_proto_order_side,
    get_proto_order_side_str,
)
from coin.support.exchange_status.logic.checker_util import (ExchangeTestRequest,
                                                             response_time_in_nano)

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)


class OrderClearError(RuntimeError):
  pass


class InvalidOrderState(RuntimeError):
  pass


class ExchangeCheckError(RuntimeError):
  def __init__(self, exchange, msg):
    self.exchange = exchange
    self.error_msg = msg


class RestStatusChecker(object):
  def __init__(self, test_request: ExchangeTestRequest):
    account_request = test_request.account_request
    key_file = get_key_file(account_request.market_type,
                            account_request.exchange,
                            account_request.owner,
                            account_request.api_version)

    self._test_request = test_request
    self._public_client = generate_public_client(account_request.market_type,
                                                 account_request.exchange,
                                                 account_request.api_version)
    self._private_client = generate_private_client(account_request.market_type,
                                                   account_request.exchange,
                                                   account_request.api_version,
                                                   key_file)
    self._product = self._private_client.ProductType.FromStr(self._test_request.product)

  def ensure_open_orders_are_cleared(self, max_num_retries=10):
    # Clear all open orders.
    for _ in range(max_num_retries):
      if self.is_open_order_cleared():
        return
      else:
        try:
          self._private_client.query_cancel_product(self._product)
        except Exception as e:
          logger.error('query_cancel_product failed: %s', e)
          raise
        time.sleep(5)  # Is 5 seconds enough
    raise OrderClearError('Failed to clear all open orders!')

  def is_order_open(self, order_element):
    logger.debug('is_order_open: private_client.query_list_orders called.')
    try:
      update = self._private_client.query_list_orders(self._product)
    except Exception as e:
      logger.error('query_list_orders failed: %s', e)
      raise

    open_orders_bundle = update.msg
    return any(
        [order_element.order_id == order.order_id for order in open_orders_bundle.each_order])

  def is_open_order_cleared(self):
    logger.debug('is_open_order_cleared: private_client.query_list_orders called.')
    try:
      update = self._private_client.query_list_orders(self._product)
    except Exception as e:
      logger.error('query_list_orders failed: %s', e)
      raise

    open_orders_bundle = update.msg
    return len(open_orders_bundle.each_order) == 0

  def is_order_filled(self, order_element):
    logger.debug('is_order_filled: private_client.query_fills called.')
    try:
      update = self._private_client.query_fills(self._product)
    except Exception as e:
      logger.error('query_fills failed: %s', e)
      raise

    fill_bundle = update.msg
    return any([order_element.order_id == fill.order_id for fill in fill_bundle.each_fill])

  def is_balance_available(self):
    logger.debug('is_balance_available: private_client.query_account_balance called.')
    try:
      update = self._private_client.query_account_balance()
    except Exception as e:
      logger.error('query_account_balance failed: %s', e)
      raise

    account_balance = update.msg
    return len(account_balance.each_balance) > 0

  def is_ticker_available(self):
    logger.debug('is_ticker_available: public_client.query_ticker called.')
    try:
      update = self._public_client.query_ticker(self._product)
    except Exception as e:
      logger.error('query_ticker failed: %s' % e)
      raise

    ticker = update.msg
    return ticker.ask > ticker.bid > 0 and ticker.last > 0

  def submit_order(self, order_element):
    order_side = get_str_proto_order_side(order_element.side)
    logger.debug('submit_order: private_client.query_submit called.')
    try:
      submit_update = self._private_client.query_submit(self._product,
                                                        order_side,
                                                        order_element.price,
                                                        order_element.qty)
    except Exception as e:
      logger.error('submit failed: %s', e)
      raise
    return submit_update

  def cancel_order(self, order_id):
    logger.debug('cancel_order: private_client.query_cancel called.')
    try:
      cancel_update = self._private_client.query_cancel(self._product, order_id)
    except Exception as e:
      logger.error('cancel %s failed %s', order_id, e)
      raise
    return cancel_update

  def gen_test_orders(self):
    logger.debug('gen_test_orders: public_client.query_ticker called.')
    try:
      update = self._public_client.query_ticker(self._product)
    except Exception as e:
      logger.error('query ticker in gen_test_orders failed: %s', e)
      raise

    ticker = update.msg
    orders = self._test_request.gen_test_orders(ticker)

    order_list = []
    for order in orders:
      order_element = ProductOrderElement(
          symbol=self._product.symbol,
          side=get_proto_order_side_str(order['side']),
          price=order['price'],
          qty=order['qty'],
          created_timestamp=get_timestamp(),
      )
      order_list.append(order_element)
    return order_list

  def test_order_api(self, order_element):
    order_attempt = OrderAttemptProto()

    # Submit test order
    try:
      submit_update = self.submit_order(order_element)
    except Exception as e:
      logger.error('test_order_api: submit order error: %s', e)
      return order_element, order_attempt

    if submit_update.msg.success:
      order_element.order_id = submit_update.msg.order_id
      order_attempt.submit_success = True
    else:
      order_attempt.submit_success = False

    order_attempt.submit_delay = response_time_in_nano(submit_update)

    if not order_attempt.submit_success:
      return order_element, order_attempt

    # Wait 2 seconds for the order to be filled or find the order in the
    # open order list,
    time.sleep(2)

    # If order can be found by query open order list, cancel it.
    try:
      is_order_open = self.is_order_open(order_element)
    except Exception as e:
      logger.error('Failed to verify that order %s is indeed open. %s', order_element.order_id, e)
      return order_element, order_attempt

    if is_order_open:
      logger.info('Order %s is open.' % order_element.order_id)
      try:
        cancel_update = self.cancel_order(submit_update.msg.order_id)
        logger.info('Cancel order %s request sent.' % order_element.order_id)
      except Exception as e:
        logger.error('test_order_api: cancel order error: %s', e)
        return order_element, order_attempt

      order_attempt.cancel_delay = response_time_in_nano(cancel_update)

      if cancel_update.msg.success:
        order_attempt.cancel_success = True
        time.sleep(2)
        try:
          if self.is_open_order_cleared():
            logger.info('Order %s is cancelled.' % order_element.order_id)
            order_attempt.is_cancelled = True
          else:
            logger.error('Cancelled order %s still found in open_order_list!',
                         order_element.order_id)
            time.sleep(2)
            self.ensure_open_orders_are_cleared()
        except Exception as e:
          logger.error('Failed to verify that cancelled orders are indeed cleared. %s' % e)
    else:
      try:
        is_order_filled = self.is_order_filled(order_element)
      except Exception as e:
        logger.error('Failed to verify that order %s is filled! %s', order_element.order_id, e)
        return order_element, order_attempt

      if is_order_filled:
        logger.info('Order %s is filled.', order_element.order_id)
        order_attempt.is_filled = True
      else:
        # Should never fall into this branch.
        logger.error('Order %s cannot be found in open_order list and filled order list!'
                     % order_element.order_id)
        self.ensure_open_orders_are_cleared()
        raise InvalidOrderState()

    return order_element, order_attempt

  def check_rest_status(self, max_num_trials):
    rest_status = RestStatusProto()
    # noinspection PyBroadException
    try:
      self.ensure_open_orders_are_cleared()
      ticker_available = self.is_ticker_available()
      rest_status.ticker_available = ticker_available
      balance_available = self.is_balance_available()
      rest_status.balance_available = balance_available
      if not (ticker_available and balance_available):
        return rest_status
    except Exception:
      return rest_status

    num_trials = 0
    order_list = []
    order_attempt_list = []
    while num_trials < max_num_trials:
      num_trials += 1

      try:
        gen_order_list = self.gen_test_orders()
      except Exception as e:
        logger.error('gen_test_orders failed! %s', e)
        print(traceback.format_exc())
        return rest_status

      for order in gen_order_list:
        order, order_attempt = self.test_order_api(order)
        order_list.append(order)
        order_attempt_list.append(order_attempt)
        time.sleep(5)
      time.sleep(30)

    rest_status.orders.extend(order_list)
    rest_status.attempts.extend(order_attempt_list)
    rest_status = summarize_rest_status(rest_status)
    return rest_status


def summarize_rest_status(rest_status):
  # Only when rest_status
  assert rest_status.ticker_available and rest_status.balance_available

  num_success = len([
      order_attempt.submit_success
      for order_attempt in rest_status.attempts
      if order_attempt.submit_success
  ])
  success_percentage = num_success / len(rest_status.attempts)
  if success_percentage >= 0.79999:
    rest_status.submit_available = True

  num_success = len([
      order_attempt.cancel_success
      for order_attempt in rest_status.attempts
      if order_attempt.cancel_success
  ])
  num_cancels = len([
      order_attempt.cancel_success
      for order_attempt in rest_status.attempts
      if not order_attempt.is_filled
  ])
  success_percentage = num_success / num_cancels
  if success_percentage >= 0.79999:
    rest_status.cancel_available = True

  rest_status.order_available = \
      rest_status.submit_available and rest_status.cancel_available
  rest_status.exchange_available = rest_status.order_available

  return rest_status


def check_rest_status(test_request: ExchangeTestRequest, num_of_trials):
  checker = RestStatusChecker(test_request)
  try:
    rest_status = checker.check_rest_status(num_of_trials)
  except Exception:
    exchange = test_request.account_request.exchange
    raise ExchangeCheckError(exchange, traceback.format_exc())
  return rest_status


def ensure_orders_are_cleared(test_request: ExchangeTestRequest):
  checker = RestStatusChecker(test_request)
  checker.ensure_open_orders_are_cleared()
